# LeetCode 787. K 站中转内最便宜的航班
# https://leetcode.cn/problems/cheapest-flights-within-k-stops/description
import heapq
from math import inf
from typing import List


class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        adj = [[] for i in range(n)]
        edge_length = {}
        for flight in flights:
            start, end, price = flight
            edge_length[(start, end)] = price
            adj[start].append(end)

        def DIJKSTRA(s, t):
            distance = [[inf for i in range(k + 2)] for j in range(n)]
            distance[s][0] = 0
            q = [(0, s, 0)]
            while q:
                du, u, step = heapq.heappop(q)
                if u == t:
                    return du
                if du > distance[u][step]:
                    continue
                if step <= k:
                    for v in adj[u]:
                        alt = distance[u][step] + edge_length[(u, v)]
                        if alt < distance[v][step + 1]:
                            distance[v][step + 1] = alt
                            heapq.heappush(q, (alt, v, step + 1))

        price = DIJKSTRA(src, dst)
        result = price if price else -1
        return result


if __name__ == '__main__':
    n = 5
    flights = [[0, 1, 5], [1, 2, 5], [0, 3, 2], [3, 1, 2], [1, 4, 1], [4, 2, 1]]
    src = 0
    dst = 2
    k = 2
    # flights = [[0,1,1],[0,2,5],[1,2,1],[2,3,1],[3,4,1]]
    # dst = 4

    # n = 13
    # flights = [[11, 12, 74], [1, 8, 91], [4, 6, 13], [7, 6, 39], [5, 12, 8], [0, 12, 54], [8, 4, 32], [0, 11, 4], [4, 0, 91],
    #  [11, 7, 64], [6, 3, 88], [8, 5, 80], [11, 10, 91], [10, 0, 60], [8, 7, 92], [12, 6, 78], [6, 2, 8], [4, 3, 54],
    #  [3, 11, 76], [3, 12, 23], [11, 6, 79], [6, 12, 36], [2, 11, 100], [2, 5, 49], [7, 0, 17], [5, 8, 95], [3, 9, 98],
    #  [8, 10, 61], [2, 12, 38], [5, 7, 58], [9, 4, 37], [8, 6, 79], [9, 0, 1], [2, 3, 12], [7, 10, 7], [12, 10, 52],
    #  [7, 2, 68], [12, 2, 100], [6, 9, 53], [7, 4, 90], [0, 5, 43], [11, 2, 52], [11, 8, 50], [12, 4, 38], [7, 9, 94],
    #  [2, 7, 38], [3, 7, 88], [9, 12, 20], [12, 0, 26], [10, 5, 38], [12, 8, 50], [0, 2, 77], [11, 0, 13], [9, 10, 76],
    #  [2, 6, 67], [5, 6, 34], [9, 7, 62], [5, 3, 67]]
    # src = 10
    # dst = 1
    # k = 10
    print(Solution().findCheapestPrice(n, flights, src, dst, k))
